import asyncio
from typing import Dict, Any
import json

class FastMCPServer:
    def __init__(self, host: str = "0.0.0.0", port: int = 8000):
        self.host = host
        self.port = port
        self.clients = {}  # 客户端连接池（addr: writer）
        self.message_handlers = {}  # 消息路由表（topic: handler）

    async def start(self):
        """启动服务端"""
        server = await asyncio.start_server(
            self.handle_client, self.host, self.port
        )
        print(f"FastMCP Server 运行在 {self.host}:{self.port}")
        async with server:
            await server.serve_forever()

    async def handle_client(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        """处理客户端连接"""
        addr = writer.get_extra_info("peername")
        self.clients[addr] = writer
        print(f"客户端 {addr} 已连接")

        try:
            while True:
                # 读取消息头（假设协议头包含消息长度）
                header = await reader.read(4)
                if not header:
                    break
                msg_length = int.from_bytes(header, byteorder="big")

                # 读取消息体
                data = await reader.read(msg_length)
                message = data.decode()
                print(f"收到来自 {addr} 的消息: {message}")

                # 处理消息并返回响应
                response = await self.process_message(message)
                await self.send_response(writer, response)

        except ConnectionResetError:
            print(f"客户端 {addr} 断开连接")
        finally:
            del self.clients[addr]
            writer.close()

    async def process_message(self, message: str) -> str:
        """消息处理逻辑（可自定义路由）"""
        # 假设消息格式为 JSON，包含 topic 和 payload
        try:
            import json
            msg_data: Dict[str, Any] = json.loads(message)
            topic = msg_data.get("topic", "default")
            payload = msg_data.get("payload", {})

            # 调用注册的处理器
            handler = self.message_handlers.get(topic, self.default_handler)
            return await handler(payload)
        except Exception as e:
            return json.dumps({"error": str(e)})

    async def default_handler(self, payload: Dict) -> str:
        """默认消息处理器"""
        return json.dumps({"status": "ok", "result": "default response"})

    async def send_response(self, writer: asyncio.StreamWriter, response: str):
        """发送响应给客户端"""
        data = response.encode()
        # 发送消息头（长度）
        writer.write(len(data).to_bytes(4, byteorder="big"))
        # 发送消息体
        writer.write(data)
        await writer.drain()

    def register_handler(self, topic: str, handler):
        """注册消息处理器"""
        self.message_handlers[topic] = handler


if __name__ == "__main__":
    server = FastMCPServer(port=8000)


    # 示例：注册自定义消息处理器
    async def echo_handler(payload: Dict) -> str:
        return json.dumps({"echo": payload.get("message")})


    server.register_handler("echo", echo_handler)

    # 启动服务
    asyncio.run(server.start())